import glob
import os
import random

import cv2
import numpy as np
import pydicom
import scipy.io
import PIL
from PIL import Image
import torch
import torchvision.transforms as transforms
from torch.utils.data import Dataset

"""
2022年9月11日 嵇老板自己重新做了数据并成功训练了自己的方法。
"""
image_get_minmax = 0.0, 1.0

proj_get_minmax = 0.0, 4.0


def normalize(data, minmax):
    data_min, data_max = minmax
    data = np.clip(data, data_min, data_max)
    data = (data - data_min) / (data_max - data_min)
    data = data.astype(np.float32)
    data = data * 255.0
    data = np.transpose(np.expand_dims(data, 2), (2, 0, 1))
    return data


class InDuDoNetDataset(Dataset):
    """
    此类专门用于InDuDoNet，读取的是使用ADN中带的制作数据的代码制作的deeplesion数据
    """

    def __init__(self, config):
        len_ratio = config['len_ratio']
        self.config = config
        self.Xma_path_list = []  # 有金属尾影的图像域图
        self.Xgt_path_list = []  # 图像域的ground truth
        self.XLI_path_list = []  # 经过LI简单处理后的图像域图
        self.Sma_path_list = []  # 有金属伪影的投影域图
        self.Sgt_path_list = []  # 投影域的ground truth
        self.SLI_path_list = []  # 经过LI简单处理后的投影域图
        self.Metal_trace_list = []  # 金属轨迹正弦图
        self.data_length = 0

        self.train_mask = np.load(r"D:\gs\code\InDuDoNet\deeplesion\train\trainmask.npy")

        all_path_list = [data_dir for data_dir in glob.glob(os.path.join(config['dataroot'], "*")) if
                         os.path.isdir(data_dir) and data_dir.find("7") != -1]

        for part_path in all_path_list:
            part_path = os.path.join(part_path, "train")
            for sample_path in glob.glob(os.path.join(part_path, "*")):
                for patient_path in glob.glob(os.path.join(sample_path, "*")):
                    # 有金属尾影的图像域图
                    tmp_list = sorted(glob.glob(os.path.join(patient_path, "ma_CT_*")))
                    self.Xma_path_list += tmp_list
                    # 图像域的ground truth
                    self.Xgt_path_list += sorted(
                        [os.path.join(patient_path, "groundtruth.mat") for i in range(len(tmp_list))])
                    # 经过LI简单处理后的图像域图
                    self.XLI_path_list += sorted(glob.glob(os.path.join(patient_path, "LI_CT_*")))
                    # 有金属伪影的投影域图
                    tmp_list = sorted(glob.glob(os.path.join(patient_path, "ma_sinogram_*")))
                    self.Sma_path_list += tmp_list
                    # 投影域的ground truth
                    tmp_list = sorted(
                        [os.path.join(patient_path, "groundtruth_sinogram.mat") for i in range(len(tmp_list))])
                    self.Sgt_path_list += tmp_list
                    # Sgt_path_list += tmp_list  # 加两次是因为有两种伪影
                    # 经过LI简单处理后的投影域图
                    tmp_list = sorted(glob.glob(os.path.join(patient_path, "LI_sinogram_*")))
                    self.SLI_path_list += tmp_list
                    # 金属轨迹正弦图
                    tmp_list = sorted(glob.glob(os.path.join(patient_path, "metal_trace_*")))
                    self.Metal_trace_list += tmp_list

        self.Xma_path_list = self.Xma_path_list[0:int(len_ratio * len(self.Xma_path_list))]
        self.Xgt_path_list = self.Xgt_path_list[0:int(len_ratio * len(self.Xgt_path_list))]
        self.XLI_path_list = self.XLI_path_list[0:int(len_ratio * len(self.XLI_path_list))]
        self.Sma_path_list = self.Sma_path_list[0:int(len_ratio * len(self.Sma_path_list))]
        self.Sgt_path_list = self.Sgt_path_list[0:int(len_ratio * len(self.Sgt_path_list))]
        self.SLI_path_list = self.SLI_path_list[0:int(len_ratio * len(self.SLI_path_list))]
        self.Metal_trace_list = self.Metal_trace_list[0:int(len_ratio * len(self.Metal_trace_list))]

        assert len(self.Xma_path_list) == len(self.Xgt_path_list) == len(self.XLI_path_list) == len(
            self.Sma_path_list) == len(self.Sgt_path_list) == len(self.SLI_path_list) == len(self.Metal_trace_list)
        self.data_length = len(self.Xma_path_list)

    def __len__(self):
        return self.data_length

    def __getitem__(self, index):
        random_mask = random.randint(0, 89)

        Xma_path = self.Xma_path_list[index]
        Xgt_path = self.Xgt_path_list[index]
        XLI_path = self.XLI_path_list[index]
        Sma_path = self.Sma_path_list[index]
        Sgt_path = self.Sgt_path_list[index]
        SLI_path = self.SLI_path_list[index]
        Metal_trace_path = self.Metal_trace_list[index]

        # ToDo 下面变量值域的原因还没有确定，是否需要归一化也没有确定
        Xma = scipy.io.loadmat(Xma_path)['image']  #
        Xgt = scipy.io.loadmat(Xgt_path)['gt_CT']  #
        XLI = scipy.io.loadmat(XLI_path)['image']  #
        Sma = scipy.io.loadmat(Sma_path)['image']  #
        Sgt = scipy.io.loadmat(Sgt_path)['gt_sinogram'].T  # 转置是因为在制作数据时没有旋转，还不确定哪种方式更好。
        SLI = scipy.io.loadmat(SLI_path)['image']  #
        Tr = scipy.io.loadmat(Metal_trace_path)['image']

        M512 = self.train_mask[:, :, random_mask]
        M = np.array(Image.fromarray(M512).resize(Xma.shape, PIL.Image.BILINEAR))
        # cv2.imshow("M", M)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()

        Mask = M.astype(np.float32)
        Mask = np.transpose(np.expand_dims(Mask, 2), (2, 0, 1))

        Xma = normalize(Xma, image_get_minmax)  # Normalization function will expand the numpy array
        Xgt = normalize(Xgt, image_get_minmax)
        XLI = normalize(XLI, image_get_minmax)
        Sma = normalize(Sma, proj_get_minmax)
        Sgt = normalize(Sgt, proj_get_minmax)
        SLI = normalize(SLI, proj_get_minmax)

        Mask = np.expand_dims(Mask, 0)
        Tr = np.expand_dims(Tr, 0)

        # The order is Xma, XLI, Xgt, mask, Sma, SLI, Sgt, Tr
        return torch.Tensor(Xma), torch.Tensor(XLI), torch.Tensor(Xgt), torch.Tensor(Mask), \
               torch.Tensor(Sma), torch.Tensor(SLI), torch.Tensor(Sgt), torch.Tensor(Tr), Xma_path


class ValDatasetDeeplesion(Dataset):
    def __init__(self, config):
        len_ratio = 0.5
        self.config = config
        self.Xma_path_list = []  # 有金属尾影的图像域图
        self.Xgt_path_list = []  # 图像域的ground truth
        self.XLI_path_list = []  # 经过LI简单处理后的图像域图
        self.Metal_trace_list = []  # 金属轨迹正弦图
        self.data_length = 0

        self.train_mask = np.load(r"D:\gs\code\InDuDoNet\deeplesion\train\trainmask.npy")

        all_path_list = [data_dir for data_dir in
                         glob.glob(os.path.join("F:\metal_artifact_data\deeplesion\original_data", "*")) if
                         os.path.isdir(data_dir) and data_dir.find("2") != -1]

        for part_path in all_path_list:
            part_path = os.path.join(part_path, "train")
            for sample_path in glob.glob(os.path.join(part_path, "*")):
                for patient_path in glob.glob(os.path.join(sample_path, "*")):
                    # 有金属尾影的图像域图
                    tmp_list = sorted(glob.glob(os.path.join(patient_path, "ma_CT_*")))
                    self.Xma_path_list += tmp_list
                    # 图像域的ground truth
                    self.Xgt_path_list += sorted(
                        [os.path.join(patient_path, "groundtruth.mat") for i in range(len(tmp_list))])
                    # 经过LI简单处理后的图像域图
                    self.XLI_path_list += sorted(glob.glob(os.path.join(patient_path, "LI_CT_*")))
                    # 金属轨迹图
                    tmp_list = sorted(glob.glob(os.path.join(patient_path, "metal_trace_*")))
                    self.Metal_trace_list += tmp_list
        self.Xma_path_list = self.Xma_path_list[0:int(len_ratio * len(self.Xma_path_list))]
        self.Xgt_path_list = self.Xgt_path_list[0:int(len_ratio * len(self.Xgt_path_list))]
        self.XLI_path_list = self.XLI_path_list[0:int(len_ratio * len(self.XLI_path_list))]
        self.Metal_trace_list = self.Metal_trace_list[0:int(len_ratio * len(self.Metal_trace_list))]

        assert len(self.Xma_path_list) == len(self.Xgt_path_list) == len(self.XLI_path_list) == len(
            self.Metal_trace_list)
        self.data_length = len(self.Xma_path_list)

    def __len__(self):
        return self.data_length

    def __getitem__(self, index):
        return_batch = dict()
        random_mask = random.randint(0, 89)

        Xma_path = self.Xma_path_list[index]
        Xgt_path = self.Xgt_path_list[index]
        XLI_path = self.XLI_path_list[index]
        Metal_trace_path = self.Metal_trace_list[index]

        # ToDo 下面变量值域的原因还没有确定，是否需要归一化也没有确定
        Xma = scipy.io.loadmat(Xma_path)['image']  #
        Xgt = scipy.io.loadmat(Xgt_path)['gt_CT']  #
        XLI = scipy.io.loadmat(XLI_path)['image']  #
        Tr = scipy.io.loadmat(Metal_trace_path)['image']

        M512 = self.train_mask[:, :, random_mask]
        M = np.array(Image.fromarray(M512).resize(Xma.shape, PIL.Image.BILINEAR))
        # cv2.imshow("M", M)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()

        Mask = M.astype(np.float32)
        Mask = np.transpose(np.expand_dims(Mask, 2), (2, 0, 1))

        Xma = normalize(Xma, image_get_minmax)  # Normalization function will expand the numpy array
        Xgt = normalize(Xgt, image_get_minmax)
        XLI = normalize(XLI, image_get_minmax)

        Mask = np.expand_dims(Mask, 0)
        Tr = np.expand_dims(Tr, 0)

        return_batch["Xma"] = torch.Tensor(Xma)
        return_batch["XLI"] = torch.Tensor(XLI)
        return_batch["Xgt"] = torch.Tensor(Xgt)
        return_batch["Tr"] = torch.Tensor(Tr)
        return_batch["Xma_path"] = Xma_path
        # The order is Xma, XLI, Xgt, mask, Sma, SLI, Sgt, Tr
        return return_batch
